
#include "Crypt.h"
#include <stdio.h>

#define	ENCRYPT_ALGORITHM CALG_RC2
#define	EXTRA_BUFFER_LENGTH 16

CCrypt::CCrypt()
{
	m_hCryptProv	= 0;
	m_hKey			= 0;
	m_hXchgKey		= 0;
	m_hHash			= 0;

	m_szOldPassword = NULL;
	m_pbKeyBlob		= NULL;
	m_pbBuffer		= NULL;

	if(!CryptAcquireContext(&m_hCryptProv, NULL, NULL, PROV_RSA_FULL, 0)){
		printf("Error taking crypt context.");
		m_hCryptProv = NULL;
		throw CCryptError();
	}
}

CCrypt::~CCrypt()
{
	DestroyCryptSession();

	if(m_hCryptProv)	{CryptReleaseContext(m_hCryptProv,0);	m_hCryptProv= 0;	}
}


/////////////////////////////////////////////////////////////////////////////
// CCrypt member functions
void CCrypt::DestroyCryptSession()
{
	if(m_pbKeyBlob)		{delete[]	m_pbKeyBlob;		m_pbKeyBlob		= NULL;	}
	if(m_pbBuffer)		{delete[]	m_pbBuffer;			m_pbBuffer		= NULL;	}
	if(m_szOldPassword)	{delete[]	m_szOldPassword;	m_szOldPassword = NULL;	}

	if(m_hKey)			{CryptDestroyKey(m_hKey);				m_hKey		= 0;	}
	if(m_hXchgKey)		{CryptDestroyKey(m_hXchgKey);			m_hXchgKey	= 0;	}
	if(m_hHash)			{CryptDestroyHash(m_hHash);				m_hHash		= 0;	}
}

BOOL CCrypt::Encrypt(void *lpBuf, int& nBufLen, const char *szPassword)
{
	int result;

	if(!szPassword)
	{
//	No password was passed
//	Encrypt the message with random key and send the key with message to a client
		DestroyCryptSession();

		if(!CryptGenKey(m_hCryptProv, ENCRYPT_ALGORITHM, CRYPT_EXPORTABLE, &m_hKey))
		{
			printf("Error generating session key.");
			return FALSE;
		}

		if(!CryptGetUserKey(m_hCryptProv, AT_KEYEXCHANGE, &m_hXchgKey))
		{
			printf("Error getting user key.");
			return FALSE;
		}

		if(!CryptExportKey(m_hKey, m_hXchgKey, SIMPLEBLOB, 0, NULL, &m_dwKeyBlobLen))
		{
			printf("Error computing blob length.");
			return FALSE;
		}

		m_pbKeyBlob = new BYTE[m_dwKeyBlobLen];

		if(!CryptExportKey(m_hKey, m_hXchgKey, SIMPLEBLOB, 0, m_pbKeyBlob, &m_dwKeyBlobLen))
		{
			printf("Error export blob key.");
			return FALSE;
		}

		m_dwBufferLen = nBufLen + EXTRA_BUFFER_LENGTH;
		m_dwCount = nBufLen;

		m_pbBuffer = new BYTE[m_dwBufferLen + m_dwKeyBlobLen + sizeof(int)];

		*((int*)m_pbBuffer) = m_dwKeyBlobLen;

		memcpy(m_pbBuffer + sizeof(int), m_pbKeyBlob, m_dwKeyBlobLen);
		memcpy(m_pbBuffer + sizeof(int) + m_dwKeyBlobLen, lpBuf, m_dwCount);

		if(!CryptEncrypt(m_hKey, 0, TRUE, 0, m_pbBuffer + m_dwKeyBlobLen + sizeof(int), &m_dwCount, m_dwBufferLen))
		{
			printf("Error encrypting message.");
			return FALSE;
		}
		
		result = TRUE;

		memcpy(lpBuf, m_pbBuffer, m_dwCount + m_dwKeyBlobLen + sizeof(int));
		nBufLen = m_dwCount + m_dwKeyBlobLen + sizeof(int);
		
		DestroyCryptSession();
	}
	else
	{
//	The message will be encrypted with a session key derived from a password
		BOOL bNeedNewKey = FALSE;
		if(m_szOldPassword)
		{
			if(strcmp(m_szOldPassword, szPassword))
			{
//	Old password and new password are different
				bNeedNewKey = TRUE;
			}
		}
		else
		{
//	Old password is NULL
			bNeedNewKey = TRUE;
		}

		if(bNeedNewKey)
		{
			DestroyCryptSession();
			
			if(!CryptCreateHash(m_hCryptProv, CALG_MD5, 0, 0, &m_hHash))
			{
				printf("Error creating Hash");
				return FALSE;
			}
			if(!CryptHashData(m_hHash, (BYTE*)szPassword, strlen(szPassword), 0))
			{
				printf("Error hashing the password");
				return FALSE;
			}
			if(!CryptDeriveKey(m_hCryptProv, ENCRYPT_ALGORITHM, m_hHash, 0, &m_hKey))
			{
				printf("Error deriving the key.");
				return FALSE;
			}
			
			CryptDestroyHash(m_hHash); m_hHash = 0;
		}

		m_dwBufferLen = nBufLen + EXTRA_BUFFER_LENGTH;
		m_dwCount = nBufLen;

		if(m_pbBuffer)	{		delete[] m_pbBuffer;	}
		
		m_pbBuffer = new BYTE[m_dwBufferLen];

		memcpy(m_pbBuffer, lpBuf, m_dwCount);

		if(!CryptEncrypt(m_hKey, 0, TRUE, 0, m_pbBuffer, &m_dwCount, m_dwBufferLen))
		{
			printf("Error encrypting message.");
			return FALSE;
		}
		
		result = TRUE;

		memcpy(lpBuf, m_pbBuffer, m_dwCount);
		nBufLen = m_dwCount;

		if(m_szOldPassword)
		{
			if(strlen(m_szOldPassword) != strlen(szPassword))
			{
				delete[] m_szOldPassword;
				m_szOldPassword = new char[strlen(szPassword)+1];
			}
		}
		else
		{
			m_szOldPassword = new char[strlen(szPassword)+1];
		}
		strcpy(m_szOldPassword, szPassword);
	}
	return result;
}

BOOL CCrypt::Decrypt(void *lpBuf, int& nBufLen, const char *szPassword)
{
	int result = TRUE;

	if(!szPassword)
	{
//	No password was passed
//	Receive the random key from message
		DestroyCryptSession();

		m_dwKeyBlobLen = *((int*)lpBuf);

		m_pbKeyBlob = new BYTE[m_dwKeyBlobLen];

		memcpy(m_pbKeyBlob, (BYTE*)lpBuf + sizeof(int), m_dwKeyBlobLen);

		if(!CryptImportKey(m_hCryptProv, m_pbKeyBlob, m_dwKeyBlobLen, 0, 0, &m_hKey))
		{
			printf("Error importing the key.");
			return FALSE;
		}
		
		m_dwBufferLen = nBufLen + EXTRA_BUFFER_LENGTH - sizeof(int) - m_dwKeyBlobLen;
		m_dwCount = nBufLen - sizeof(int) - m_dwKeyBlobLen;

		m_pbBuffer = new BYTE[m_dwBufferLen];

		memcpy(m_pbBuffer, (BYTE*)lpBuf + m_dwKeyBlobLen + sizeof(int), m_dwCount);

		if(!CryptDecrypt(m_hKey, 0, TRUE, 0, m_pbBuffer, &m_dwCount))
		{
			printf("Error decrypting message.");
			return FALSE;
		}
		
		memcpy(lpBuf, m_pbBuffer, m_dwCount);
		nBufLen = m_dwCount;
		
		DestroyCryptSession();
	}
	else
	{
//	The message will be decrypted with a session key derived from a password
		BOOL bNeedNewKey = FALSE;
		if(m_szOldPassword)
		{
			if(strcmp(m_szOldPassword, szPassword))
			{
//	Old password and new password are different
				bNeedNewKey = TRUE;
			}
		}
		else
		{
//	Old password is NULL
			bNeedNewKey = TRUE;
		}

		if(bNeedNewKey)
		{
			DestroyCryptSession();
			
			if(!CryptCreateHash(m_hCryptProv, CALG_MD5, 0, 0, &m_hHash))
			{
				printf("Error creating Hash");
				return FALSE;
			}
			if(!CryptHashData(m_hHash, (BYTE*)szPassword, strlen(szPassword), 0))
			{
				printf("Error hashing the password");
				return FALSE;
			}
			if(!CryptDeriveKey(m_hCryptProv, ENCRYPT_ALGORITHM, m_hHash, 0, &m_hKey))
			{
				printf("Error deriving the key.");
				return FALSE;
			}
			
			CryptDestroyHash(m_hHash); m_hHash = 0;
		}

		m_dwBufferLen = nBufLen + EXTRA_BUFFER_LENGTH;
		m_dwCount = nBufLen;

		if(m_pbBuffer)	{		delete[] m_pbBuffer;	}

		m_pbBuffer = new BYTE[m_dwBufferLen];

		memcpy(m_pbBuffer, lpBuf, m_dwCount);

		if(!CryptDecrypt(m_hKey, 0, TRUE, 0, m_pbBuffer, &m_dwCount))
		{
			printf("Error decrypting message.");
			return FALSE;
		}
		
		nBufLen = m_dwCount;
		memcpy(lpBuf, m_pbBuffer, m_dwCount);

		if(m_szOldPassword)
		{
			if(strlen(m_szOldPassword) != strlen(szPassword))
			{
				delete[] m_szOldPassword;
				m_szOldPassword = new char[strlen(szPassword)+1];
			}
		}
		else
		{
			m_szOldPassword = new char[strlen(szPassword)+1];
		}
		strcpy(m_szOldPassword, szPassword);
	}
	return result;
}
